from __future__ import print_function, division
import time
#Importing relevant modules and packages
import numpy as np
import matplotlib.pyplot as plt
import random
from RichardsModel_1D import *
from Parameters import *
from scipy import integrate
from tqdm import tqdm



p=Loam() #Soil parameters

#Space and time parameters
Nr,Nt,Nz,dr,dt,dz,Np,Na,Np,N_obs,Hz=circular_parameters()
DeltaT,Nsim=time_parameters()

N = 50 # Number of particles

#Irrigation amount
# uu=np.zeros(360*40)
# for i in range(360*40):
#     if i < 360:
#         uu[i]=-1e-05 #Moisture is supplied only at the first time instant
#     else:
#         uu[i]=0
    # uu[0*i]=-1e-05

u=np.zeros(3*120)
uu=np.zeros(Nsim)
for i in range(3*120):
    if i < 5:
        u[i]=-1e-06 #Moisture is supplied only at the first time instant
    else:
        u[i]=0
uu=np.tile(u,(1,120))
uu=np.transpose(uu)


Kc=0.88*np.ones(Nsim)      #Crop coefficient

Et=1.389e-08*np.ones(Nsim) #Reference evapotranspiration




#Initial condition [pressure head]
x0=-1.05*np.ones(Nz)


#Numpy array to store the simulation results [pressure head]
xx=np.zeros((Nsim+1,Nz))
xx1=np.zeros((Nsim+1,Nz))
ode1=np.zeros((Nsim+1,Nz))
ode2=np.zeros((Nsim+1,Nz))
ode4=np.zeros((Nsim+1,Nz))
xx[0,:]=x0
xx1[0,:]=x0



#Model uncertainty
#Model uncertainty
www1 = 4e-9 # process noise
vv1 = 8e-7*np.ones(N_obs) # measurement noise

np.random.seed(2)
w=np.random.normal(0,np.sqrt(www1),[Nsim+1,Nz])
# w=0.01*np.random.randn(Nsim,Nz)
np.random.seed(3)
vv=np.random.randn(Nsim+1,N_obs)*vv1 #measurement noise

#Numpy array to store the simulation results [volumetric moisture content]
yy=np.zeros((Nsim+1,N_obs))

yy[0,:]=getOutputs1D(x0).full().ravel()+vv[0,:]


#unknown inputs
a1=0*np.ones((Nsim+1,Nz))
# a=0.000009*np.ones((Nsim+1,Nz))#constant UIs
a=0.00003*np.ones((Nsim+1,Nz))
# for i in range(Nsim+1):
#     a[i,:]=a[i,:]*np.cos(np.pi/(Nsim+1)*i)
# a_est=0.000006*np.ones((Nsim+1,Nz))




    
for i in tqdm(range(1, Nsim+1)):
    xx[i,:]= DM(ODE_discrete(tuple(xx[i-1]), uu[i-1],a[i-1],DeltaT)).full().ravel()+w[i-1]    
    yy[i,:]=getOutputs1D(xx[i,:]).full().ravel()+vv[i,:]

# for i in tqdm(range(1, Nsim+1)):
#     # Data with noise without UIs
#     ode2[i-1,:]= RichardsModel_1D(tuple(xx1[i-1]), uu[i-1],ww[i-1],a1[i-1],Kc[i-1],Et[i-1],p)
#     xx1[i,:]=xx1[i-1,:]+DeltaT*ode2[i-1,:]
#     yy1[i,:]=getOutputs1D_allnodes(xx1[i,:],p).ravel()+vv[i,:]
    


t= list(range(Nsim+1))
plt.figure(1)
for i in range (Nz):
    plt.subplot(4, 4, i+1)
    plt.plot(t,xx[:,i],'r')
plt.show()
